/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
 * Description: broadcast nhr v2.0 算法操作类实现文件
 */

#include "broadcast_nhr.h"
#include "scatter_nhr.h"
#include "all_gather_nhr.h"
#include "broadcast_nhr_oneshot.h"

namespace hccl {
BroadcastNHR::BroadcastNHR(const HcclDispatcher dispatcher)
    : NHRBase(dispatcher)
{
}

BroadcastNHR::~BroadcastNHR()
{
}

HcclResult BroadcastNHR::RunAsync(const u32 rank, const u32 rankSize,
    const std::vector<std::shared_ptr<Transport> > &links)
{
    CHK_SMART_PTR_NULL(dispatcher_);
    CHK_PTR_NULL(stream_.ptr());
    HCCL_INFO("[BroadcastNHR][RunAsync] run: rank[%u] totalrank[%u] count[%llu]", rank, rankSize, count_);

    if (rankSize == 1) {
        return HCCL_SUCCESS;
    }

    CHK_PRT_RET(links.size() < rankSize,
        HCCL_ERROR("[BroadcastNHR][RunAsync] rank[%u] linksize[%llu] is less than rank size", rank, links.size()),
        HCCL_E_INTERNAL);

    u32 unitSize = DataUnitSize(dataType_);
    CHK_PRT_RET(unitSize == 0, HCCL_ERROR("[BroadcastNHR][RunAsync] rank[%u] unit data size is zero", rank),
        HCCL_E_INTERNAL);

    DeviceMem srcMem = inputMem_.range(baseOffset_, count_ * unitSize);
    DeviceMem dstMem = outputMem_.range(baseOffset_, count_ * unitSize);

    if (inputMem_ != outputMem_ && rank == root_) {
        CHK_RET(HcclD2DMemcpyAsync(dispatcher_, dstMem, srcMem, stream_));
    }

    // 准备slice
    PrepareSlice(rank, rankSize);

    // scatter
    ScatterNHR tempAlgScatter(dispatcher_);
    tempAlgScatter.CloseBarrier();
    CHK_RET(tempAlgScatter.Prepare(srcMem, srcMem, srcMem, count_, dataType_, stream_,
        reductionOp_, root_, slices_, baseOffset_));

    CHK_RET(tempAlgScatter.RegisterProfiler(
        profilerInput_.planeID, profilerInput_.stage, profilerInput_.step, stream_));
    CHK_RET(tempAlgScatter.RunAsyncWithReorder(rank, rankSize, links));

    // allgather
    AllGatherNHR tempAlgAllgather(dispatcher_);
    CHK_RET(tempAlgAllgather.Prepare(srcMem, srcMem, srcMem, count_, dataType_, stream_,
        reductionOp_, root_, slices_, baseOffset_));

    CHK_RET(tempAlgAllgather.RegisterProfiler(
        profilerInput_.planeID, profilerInput_.stage, profilerInput_.step, stream_));
    CHK_RET(tempAlgAllgather.RunAsyncWithReorder(rank, rankSize, links));

    HCCL_INFO("[BroadcastNHR][RunAsync] finished: rank[%u] end count[%llu]", rank, count_);
    return HCCL_SUCCESS;
}

HcclResult BroadcastNHR::PrepareSlice(const u32 rank, const u32 rankSize)
{
    u32 unitSize = DataUnitSize(dataType_);

    //  所有的数据平均到每个rank上
    u64 sizeAvg = ((count_ + rankSize - 1) / rankSize) * unitSize;
    u64 sizePerSlice = AlgTemplateBase::RoundUpWithDivisor(sizeAvg, HCCL_MIN_SLICE_ALIGN);
    HCCL_DEBUG("[BroadcastNHR][PrepareSlice] bcast total count[%llu] sizeAverage[%llu] sizePerSlice after aligns[%llu]",
        count_, sizeAvg, sizePerSlice);

    // 准备slice
    slices_.resize(rankSize);
    u64 sizeResidue = count_ * unitSize;
    u64 sizePerRound = 0;

    for (u32 i = 0; i < rankSize; i++) {
        sizePerRound = (sizeResidue > sizePerSlice) ? sizePerSlice : sizeResidue;
        slices_[i].offset = count_ * unitSize - sizeResidue;
        slices_[i].size = sizePerRound;

        sizeResidue -= sizePerRound;
        HCCL_DEBUG("[BroadcastNHR][PrepareSlice] rank[%u] default slice[%u]: offset: [%llu] size[%llu]", rank, i,
            slices_[i].offset, slices_[i].size);
    }
    return HCCL_SUCCESS;
}

}  // namespace hccl
